package dfxnftables

import (
	"fmt"
	"log"
	"os"
	"os/exec"

	"github.com/google/nftables"
)

const (
	NFT_ADDRESS_FAMILY_IPV4   = "ip"
	NFT_ADDRESS_FAMILY_IPV6   = "ip6"
	NFT_ADDRESS_FAMILY_NET    = "inet"
	NFT_ADDRESS_FAMILY_ARP    = "arp"
	NFT_ADDRESS_FAMILY_BRIDGE = "bridge"
)
// NfMeta 统一的元信息结构
// 用于聚合常用且重复的字段，避免在 Set/Chain/Rule/Table 等结构体中反复声明
// - TableName: 所属表名
// - AddressFamily: 地址族，取值参见常量，如 ip/ip6/inet
type NfMeta struct {
	TableName     string
	AddressFamily string
}

// Family 返回 nftables 的地址族
// 返回值为 nftables.TableFamily 类型，表示当前地址族对应的 nftables 表族
func (m NfMeta) Family() nftables.TableFamily {
	return familyFromString(m.AddressFamily)
}

// Table 返回带有表名与地址族的 *nftables.Table
// 返回值为 *nftables.Table 类型，包含表名和地址族信息
func (m NfMeta) Table() *nftables.Table {
	return &nftables.Table{Name: m.TableName, Family: m.Family()}
}

// Validate 校验必填字段
// 检查 TableName 是否为空以及 AddressFamily 是否为支持的地址族
// 返回 error 类型，如果校验失败则返回错误信息
func (m NfMeta) Validate() error {
	if m.TableName == "" {
		return fmt.Errorf("TableName 不能为空")
	}
	switch m.AddressFamily {
	case NFT_ADDRESS_FAMILY_IPV4, NFT_ADDRESS_FAMILY_IPV6, NFT_ADDRESS_FAMILY_NET, NFT_ADDRESS_FAMILY_ARP, NFT_ADDRESS_FAMILY_BRIDGE:
		return nil
	default:
		return fmt.Errorf("不支持的 AddressFamily: %s", m.AddressFamily)
	}
}

// ResolveTable 尝试从 nftables.Conn 中获取带有内核句柄的 Table
// 仅当表已存在于内核时才会返回成功，否则返回错误
// 参数 c: *nftables.Conn 类型，表示与 nftables 内核模块的连接
// 返回值: *nftables.Table 类型表示找到的表，error 类型表示可能发生的错误
func (m NfMeta) ResolveTable(c *nftables.Conn) (*nftables.Table, error) {
	if err := m.Validate(); err != nil {
		return nil, err
	}
	tables, err := c.ListTables()
	if err != nil {
		return nil, fmt.Errorf("list tables failed: %w", err)
	}
	fam := m.Family()
	for _, tbl := range tables {
		if tbl.Name == m.TableName && tbl.Family == fam {
			return tbl, nil
		}
	}
	return nil, fmt.Errorf("table %s (%s) not found", m.TableName, m.AddressFamily)
}

// familyFromString 将字符串的地址族映射到 nftables.TableFamily
// 参数 s: string 类型，表示地址族字符串
// 返回值: nftables.TableFamily 类型，表示对应的 nftables 表族
func familyFromString(s string) nftables.TableFamily {
	switch s {
	case NFT_ADDRESS_FAMILY_IPV4:
		return nftables.TableFamilyIPv4
	case NFT_ADDRESS_FAMILY_IPV6:
		return nftables.TableFamilyIPv6
	case NFT_ADDRESS_FAMILY_NET:
		return nftables.TableFamilyINet
	case NFT_ADDRESS_FAMILY_ARP:
		return nftables.TableFamilyARP
	case NFT_ADDRESS_FAMILY_BRIDGE:
		return nftables.TableFamilyBridge
	default:
		return nftables.TableFamilyINet
	}
}

// RequireNftables 检查运行环境是否满足 nftables 使用要求
// 包括是否以 root 用户身份运行、是否加载了 nf_tables 内核模块
// 如果不满足条件，程序将输出错误信息并退出
func RequireNftables() {
	if os.Geteuid() != 0 {
		log.Fatalln("需要root用户执行")
	}
	if _, err := os.Stat("/proc/net/netfilter/nfnetlink"); err != nil {
		if _, err2 := exec.Command("modprobe", "nf_tables").CombinedOutput(); err2 != nil {
			log.Fatalln("请安装nftables")
		}
	}
}